package org.jberney.wvdexporter.io;

import java.awt.image.BufferedImage;
import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.util.Deque;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.Map;
import java.util.NavigableMap;
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import org.jberney.wvdexporter.video.WvdFrame;
import org.jberney.wvdexporter.video.WvdFrameHeader;
import org.jberney.wvdexporter.video.WvdFrameTimeStamp;

import com.xuggle.mediatool.IMediaWriter;
import com.xuggle.xuggler.ICodec;

public class WvdEncoder {

	private String root;
	private Set<WvdEncodeProgressListener> listeners;
	private int inputIndex;
	private int streamId;
	private int width;
	private int height;

	public WvdEncoder(String root, int inputIndex, int streamId) {
		this.root = root;
		this.listeners = new HashSet<WvdEncodeProgressListener>();
		this.inputIndex = inputIndex;
		this.streamId = streamId;
	}

	protected long readFrameTimeStamp(InputStream inputStream)
			throws IOException {
		long timeStamp = -1L;
		if (inputStream.available() > 0) {
			byte[] timeStampBytes = new byte[WvdFrameTimeStamp.LENGTH];
			inputStream.read(timeStampBytes);
			WvdFrameTimeStamp wvdFrameTimeStamp = new WvdFrameTimeStamp(
					timeStampBytes);
			timeStamp = wvdFrameTimeStamp.getTimeStamp();
		}
		return timeStamp;
	}

	protected WvdFrameHeader readFrameHeader(InputStream inputStream)
			throws IOException {
		byte[] headerBytes = new byte[WvdFrameHeader.LENGTH];
		inputStream.read(headerBytes);
		return new WvdFrameHeader(headerBytes);
	}

	protected WvdFrame readNextFrame(InputStream inputStream,
			WvdFrameHeader wvdFrameHeader) throws IOException {
		int frameLength = wvdFrameHeader.getFrameLength();
		byte[] frameBytes = new byte[frameLength];
		inputStream.read(frameBytes);
		return new WvdFrame(wvdFrameHeader, frameBytes);
	}

	protected InputStream getInputStream(File file)
			throws FileNotFoundException {
		FileInputStream fileInputStream = new FileInputStream(file);
		return new BufferedInputStream(fileInputStream);
	}

	protected File getNextFile(InputStream inputStream) throws IOException {
		byte[] bytes = new byte[inputStream.available()];
		inputStream.read(bytes);
		String nextPath = new String(bytes);
		if (!nextPath.startsWith("/")) {
			nextPath = String.format("%s/%s", this.root, nextPath);
		}
		return new File(nextPath);
	}

	protected NavigableMap<File, Deque<Integer>> getOffsets(File file)
			throws IOException {
		NavigableMap<File, Deque<Integer>> offsets = new TreeMap<File, Deque<Integer>>();
		InputStream inputStream = this.getInputStream(file);
		int offset = 0;
		long timeStamp = -1L;
		boolean done = true;
		Deque<Integer> fileOffsets = new LinkedList<Integer>();
		do {
			timeStamp = this.readFrameTimeStamp(inputStream);
			offset += WvdFrameTimeStamp.LENGTH;
			done = timeStamp == -1L;
			if (!done) {
				if (timeStamp != WvdFrameTimeStamp.NEXTFILE_TIMESTAMP) {
					fileOffsets.add(offset - WvdFrameTimeStamp.LENGTH);
					WvdFrameHeader wvdFrameHeader = this
							.readFrameHeader(inputStream);
					if (this.width == 0 || this.height == 0) {
						this.width = wvdFrameHeader.getWidth();
						this.height = wvdFrameHeader.getHeight();
					}
					int frameLength = wvdFrameHeader.getFrameLength();
					int skipped = 0;
					while (skipped < frameLength) {
						skipped += inputStream.skip(frameLength - skipped);
					}
					offset += WvdFrameHeader.LENGTH + frameLength;
				} else {
					offsets.put(file, new LinkedList<Integer>(fileOffsets));
					file = this.getNextFile(inputStream);
					inputStream.close();
					inputStream = this.getInputStream(file);
					fileOffsets.clear();
					offset = 0;
				}
			}
		} while (!done);
		offsets.put(file, fileOffsets);
		inputStream.close();
		return offsets;
	}

	public void encode(File file, IMediaWriter mediaWriter) throws IOException,
			InterruptedException {

		NavigableMap<File, Deque<Integer>> offsets = this.getOffsets(file);

		int availableProcessors = Runtime.getRuntime().availableProcessors();
		ExecutorService executorService = Executors
				.newFixedThreadPool(availableProcessors);

		int streamIndex = mediaWriter
				.addVideoStream(this.inputIndex, this.streamId,
						ICodec.ID.CODEC_ID_H264, this.width, this.height);

		for (Map.Entry<File, Deque<Integer>> entry : offsets.entrySet()) {

			file = entry.getKey();
			Deque<Integer> fileOffsets = entry.getValue();

			InputStream[] inputStreams = new InputStream[availableProcessors];

			for (int i = 0; i < availableProcessors; i++) {
				FileInputStream fileInputStream = new FileInputStream(file);
				inputStreams[i] = new BufferedInputStream(fileInputStream);
			}

			int[] previousPeeks = new int[availableProcessors];

			while (!fileOffsets.isEmpty()) {

				int numThreads = Math.min(availableProcessors,
						fileOffsets.size());

				CountDownLatch countDownLatch = new CountDownLatch(numThreads);

				Map<Long, BufferedImage> frames = new TreeMap<Long, BufferedImage>();
				for (int i = 0; i < numThreads; i++) {

					int offset = fileOffsets.pop();
					int skip = offset - previousPeeks[i];
					if (!fileOffsets.isEmpty()) {
						previousPeeks[i] = fileOffsets.peek();
					}

					executorService.execute(new EncoderRunnable(this, frames,
							inputStreams[i], skip, countDownLatch));

				}

				countDownLatch.await();

				for (Map.Entry<Long, BufferedImage> frame : frames.entrySet()) {
					long timeStamp = frame.getKey();
					BufferedImage bufferedImage = frame.getValue();
					mediaWriter.encodeVideo(streamIndex, bufferedImage,
							timeStamp, TimeUnit.MICROSECONDS);
				}

			}

		}

		executorService.shutdown();

	}

	public void addProgressListener(WvdEncodeProgressListener listener) {
		this.listeners.add(listener);
	}

	private static class EncoderRunnable implements Runnable {

		private WvdEncoder wvdEncoder;
		private Map<Long, BufferedImage> frames;
		private InputStream inputStream;
		private CountDownLatch countDownLatch;

		public EncoderRunnable(WvdEncoder wvdEncoder,
				Map<Long, BufferedImage> frames, InputStream inputStream,
				int skip, CountDownLatch countDownLatch) throws IOException {
			this.wvdEncoder = wvdEncoder;
			this.frames = frames;
			this.inputStream = inputStream;
			int skipped = 0;
			while (skipped < skip) {
				skipped += this.inputStream.skip(skip - skipped);
			}
			this.countDownLatch = countDownLatch;
		}

		public void run() {

			try {
				long timeStamp = this.wvdEncoder
						.readFrameTimeStamp(inputStream);
				WvdFrameHeader wvdFrameHeader = this.wvdEncoder
						.readFrameHeader(inputStream);
				WvdFrame wvdFrame = this.wvdEncoder.readNextFrame(inputStream,
						wvdFrameHeader);
				BufferedImage bufferedImage = wvdFrame.getFrameAsImage();
				this.frames.put(timeStamp, bufferedImage);
			} catch (IOException e) {
				e.printStackTrace();
			}

			this.countDownLatch.countDown();

		}

	}

}
